# -*- coding:utf-8 -*-
"""
@file name  : tensorboard_methods_2.py
@author     : QuZhang
@date       : 2020-12-24 20:36
@brief      : tensorboard方法使用2
"""
from tools.common_tools import set_seed
from torch.utils.tensorboard import SummaryWriter
import torch
import time
import os
from torchvision.transforms import transforms
from tools.my_dataset import RMBDataset
from torch.utils.data import DataLoader
import torchvision.utils as vutils
from model.lenet import LeNet


set_seed(1)

# ------------- 3 add_image ---------------
# 图片可视化
# flag = True
flag = False
if flag:
    writer = SummaryWriter(comment='test_your_comment', filename_suffix='_test_your_filename_suffix')

    # img 1 random
    fake_img = torch.randn(3, 512, 512)  # 3*512*512
    writer.add_image('fake_img', fake_img, 1)
    time.sleep(1)

    # img 2 ones
    fake_img = torch.ones(3, 512, 512)  # 像素值:255，白色
    time.sleep(1)
    writer.add_image('fake_img', fake_img, 2)

    # img 3 1.1
    fake_img = torch.ones(3, 512, 512) * 1.1
    time.sleep(1)
    writer.add_image('fake_img', fake_img, 3)

    # img 4 HW
    fake_img = torch.rand(512, 512)  # 灰度
    writer.add_image('fake_img', fake_img, 4, dataformats="HW")  # 图片格式：高*宽

    # img 5 HWC
    fake_img = torch.rand(512, 512, 3)
    writer.add_image('fake_img', fake_img, 5, dataformats='HWC')  # 数据格式：高*宽*通道数

    writer.close()  # 这一行必须有

# ------------- 4 make_grid ---------------
# 将多张图片显示在不同网格中
# flag = True
flag = False
if flag:
    writer = SummaryWriter(comment='test_make_grid', filename_suffix="_test_your_filename_suffix")

     # 构建数据路径
    split_dir = os.path.join("..", "..", 'data', "rmb_split")
    train_dir = os.path.join(split_dir, 'train')

    # 图像变换
    transforms_compose = transforms.Compose([
        transforms.Resize((32, 64)),
        transforms.ToTensor(),
    ])

    # 数据加载器
    train_data = RMBDataset(data_dir=train_dir, transform=transforms_compose)
    train_loader = DataLoader(dataset=train_data, batch_size=16, shuffle=True)

    # 读取数据
    """ 
    dataloader本质是一个可迭代对象,可用iter进行访问返回一个迭代器，之后用next获取数据
    """
    data_batch, label_batch = next(iter(train_loader))  # 读取一批图片
    print("data_batch shape: ", data_batch.shape)

    # 可视化
    # img_grid = vutils.make_grid(data_batch, nrow=4, normalize=True, scale_each=True)
    # make_grid参数里的tensor要为4维张量,(可视化的图片数，通道，高，宽)
    img_grid = vutils.make_grid(data_batch, nrow=4, normalize=False, scale_each=False)  # 创建网格对象，将要可视化的图像张量添加到网格
    writer.add_image('input img', img_grid, 0)  # 将网格对象添加到可视化图像里

    writer.close()

# ------------ 5 add_graph ----------------
# 可视化计算图,用于优化和调参
flag = True
if flag:
    writer = SummaryWriter(comment='calc_graph', filename_suffix="LeNet_visualization")

    fake_img = torch.randn(1, 3, 32, 32)  # 数据
    lenet = LeNet(classes=2)  # 模型
    writer.add_graph(lenet, fake_img)
    writer.close()

    # 查看模型信息，用于调试
    from torchsummary import summary
    print(summary(lenet, (3, 32, 32), device="cpu"))
